"""


Version: 0.1
Author: lk
Date: 2022-03-08
"""
import torch
from torch.utils import data
import dgl
import networkx as nx
class GraphBatcher():
    def __init__(self,graph_q, graph_k):
        self.graph_q = graph_q
        self.graph_k = graph_k
    def __str__(self):
        return str(self.graph_q)+"\n"+str(self.graph_k)

def batcher(batch):
    graph_k,graph_q = zip(*batch)
    graph_k,graph_q = dgl.batch(graph_k),dgl.batch(graph_q)
    return GraphBatcher(graph_q,graph_k)

class GraphDataset(data.Dataset):
    def __init__(self,rw_hops=128):
        self.rw_hops = rw_hops
        graphs = []
        for i in range(5):
            graphs.append(dgl.from_networkx(nx.cycle_graph(10)))
        self.graph = dgl.batch(graphs)
    def __len__(self):
        return self.graph.number_of_nodes()
    def get_2_traces(self,item):
        traces=[]
        for i in range(2):
            traces.append(dgl.sampling.random_walk(
                self.graph,
                [item],
                length=self.rw_hops
            )[0])
        return torch.cat(traces,dim=0)

    def __getitem__(self, item):
        traces = self.get_2_traces(item)
        # return traces[0],traces[1]
        q_nodes,k_nodes=traces[0],traces[1]
        return dgl.node_subgraph(self.graph,q_nodes),dgl.node_subgraph(self.graph,k_nodes)
        # print(traces)
if __name__ == '__main__':
    coauthor = dgl.data.CoauthorCSDataset(".")
    print(coauthor)
    print(coauthor[0])
    dataset = GraphDataset(5)
    dataloader = data.DataLoader(dataset,batch_size=2,collate_fn=batcher)
    # for graph_q, graph_k in dataset:
    #     print(graph_q)
    #     print(graph_k)
    for batch in dataloader:
        print(batch)
        print(type(batch))
        break